import numpy as np
import pandas as pd
import networkx as nx
import matplotlib.pyplot as plt
from pathlib import Path
import sys
import os
from UAS.data.graph_validator import GraphValidator
from UAS.data.graph_utils import GraphUtils

sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

class GraphBuilder:
    def __init__(self, excel_parser):
        self.parser = excel_parser
        self.G = nx.DiGraph()
        self.node_mapping = {'S': 0, 'F': 1}
        self.reverse_mapping = {0: 'S', 1: 'F'}
        self._create_node_mapping()  # Initialize mapping immediately

    def build_graph(self):
        """构建图结构并返回关键数据"""
        print("Starting graph building process...")
        self._create_node_mapping()
        print("Node mapping complete, building graph...")
        self._build_graph()
        print("Graph building complete")
        # 获取验证通过的点位并转换为编码，同时获取对应的profit
        valid_points = []
        profits = []
        for p in self.parser.points:
            if p in self.node_mapping:
                valid_points.append(self.node_mapping[p])
                profits.append(int(round(self.parser.profits[p])))
        
        # 生成邻接矩阵
        n = len(self.node_mapping)
        adj_matrix = np.zeros((n, n), dtype=int)
        for u in self.parser.time_matrix:
            for v, weight in self.parser.time_matrix[u].items():
                i = self.node_mapping[u]
                j = self.node_mapping[v]
                # Handle special cases
                if u == 'F' or v == 'S':
                    adj_matrix[i][j] = 0
                else:
                    # Convert to float whether input is int or float
                    adj_matrix[i][j] = int(round(weight * 10)) if weight > 0 else 0
                
        # 返回有效点位编码、profit、邻接矩阵和节点映射
        return valid_points, profits, adj_matrix, self.node_mapping
        
    def _validate_node_mapping(self, node):
        """Ensure node exists in mapping, create if not"""
        if node not in self.node_mapping:
            idx = len(self.node_mapping)
            self.node_mapping[node] = idx
            self.reverse_mapping[idx] = node

    def _create_node_mapping(self):
        """Generate node mapping for all node types"""
        idx = 2
        for point in self.parser.points:
            # Handle all node types
            if isinstance(point, int):
                point = str(point)  # Convert integers to strings
            if point not in ['S', 'F']:
                self._validate_node_mapping(point)
                idx += 1
                
    def _build_graph(self):
        """构建网络图，处理时间矩阵"""
        # Ensure all nodes are mapped
        for u in self.parser.time_matrix:
            self._validate_node_mapping(u)
            for v in self.parser.time_matrix[u]:
                self._validate_node_mapping(v)
                
        # 转换时间并处理NaN
        time_matrix = {
            u: {
                # Multiply by 10, round and convert to float
                v: int(round(time * 10)) if not np.isnan(time) else 0
                for v, time in row.items()
            }
            for u, row in self.parser.time_matrix.items()
        }

        #print(f"Time matrix after conversion: {time_matrix}")

        # 构建图结构
        edge_count = 0
        for u in time_matrix:
            for v, weight in time_matrix[u].items():
                # 强制对角线为0
                if u == v:
                    weight = 0
                
                if weight >= 0:  # 包含0值路径但过滤负值
                    self.G.add_edge(
                        self.node_mapping[u],
                        self.node_mapping[v],
                        weight=weight,
                        label=str(weight)
                    )
                    edge_count += 1
        #print(f"Graph built with {edge_count} edges")
        #print(f"Graph nodes: {self.G.nodes()}")
        #print(f"Graph edges: {self.G.edges(data=True)}")

    def calculate_floyd(self):
        """执行Floyd算法并保存结果，同时记录路径信息"""
        # Ensure all nodes are mapped
        self._create_node_mapping()
        
        n = len(self.node_mapping)
        adj_matrix = np.zeros((n, n))
        
        # 构建邻接矩阵（0表示无连接）
        for u in self.parser.time_matrix:
            if u not in self.node_mapping:
                continue
            for v, weight in self.parser.time_matrix[u].items():
                if v not in self.node_mapping:
                    continue
                i = self.node_mapping[u]
                j = self.node_mapping[v]
                # Convert to float whether input is int or float
                adj_matrix[i][j] = float(weight) if weight > 0 else 0.0

        # 使用GraphUtils执行Floyd算法
        dist, next_node = GraphUtils.floyd_warshall(adj_matrix)
        self.next_node = next_node  # 保存路径追踪矩阵
        self._save_results(adj_matrix, dist)
        return dist

    def _reconstruct_path(self, u, v):
        """根据next矩阵重建最短路径"""
        return GraphUtils.reconstruct_path(self.next_node, u, v)
    
    def _save_results(self, adj_matrix, floyd_matrix):
        """将结果保存到Excel文件"""
        with pd.ExcelWriter('output/verified_matrix.xlsx') as writer:
            # 保存邻接矩阵
            adj_df = pd.DataFrame(adj_matrix).rename(
                columns=self.reverse_mapping,
                index=self.reverse_mapping
            )
            adj_df.to_excel(writer, sheet_name='Adjacency Matrix')
            
            # 保存Floyd结果
            floyd_df = pd.DataFrame(floyd_matrix).rename(
                columns=self.reverse_mapping,
                index=self.reverse_mapping
            )
            floyd_df.to_excel(writer, sheet_name='Floyd Matrix')
            
            # 保存节点映射
            mapping_df = pd.DataFrame.from_dict(
                self.node_mapping, 
                orient='index', 
                columns=['Code']
            )
            mapping_df.to_excel(writer, sheet_name='Node Mapping')
    
    def visualize_graph(self):
        """可视化图形并保存"""
        plt.figure(figsize=(12, 8))
        pos = nx.spring_layout(self.G)
        
        # 绘制节点和边
        nx.draw_networkx_nodes(self.G, pos, node_size=80)
        nx.draw_networkx_edges(self.G, pos, arrowstyle='->', arrowsize=20)
        
        # 添加标签
        node_labels = {n: self.reverse_mapping[n] for n in self.G.nodes()}
        edge_labels = nx.get_edge_attributes(self.G, 'label')
        nx.draw_networkx_labels(self.G, pos, labels=node_labels)
        nx.draw_networkx_edge_labels(self.G, pos, edge_labels=edge_labels)
        
        plt.title("Orienteering Graph Structure")
        plt.savefig('output/graph_visualization.png')
        plt.close()

if __name__ == '__main__':
    from UAS.data.excel_parser import ExcelParser
    parser = ExcelParser("valid_points.xlsx", "valid_matrix.xlsx")
    parser.load_data()    
    builder = GraphBuilder(parser)
    # Build the graph before validation
    builder.build_graph()
    validator = GraphValidator(builder)
    violations = validator.validate_paths()
    builder.visualize_graph()
    print(f"发现{len(violations)}条路径违规，详情请查看output/violations.log")
